{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy.optimize import fsolve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# esthetic definitions for the plots\n",
    "font = {'size'   : 16, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a function returning the three masses given the Delta m^2, the total mass, and the hierarchy (e.g. 'IN' or 'IH')\n",
    "# taken from a piece of MontePython written by Thejs Brinckmann\n",
    "def get_masses(delta_m_squared_atm, delta_m_squared_sol, sum_masses, hierarchy):\n",
    "    # any string containing letter 'n' will be considered as refering to normal hierarchy\n",
    "    if 'n' in hierarchy.lower():\n",
    "        # Normal hierarchy massive neutrinos. Calculates the individual\n",
    "        # neutrino masses from M_tot_NH and deletes M_tot_NH\n",
    "        #delta_m_squared_atm=2.45e-3\n",
    "        #delta_m_squared_sol=7.50e-5\n",
    "        m1_func = lambda m1, M_tot, d_m_sq_atm, d_m_sq_sol: M_tot**2. + 0.5*d_m_sq_sol - d_m_sq_atm + m1**2. - 2.*M_tot*m1 - 2.*M_tot*(d_m_sq_sol+m1**2.)**0.5 + 2.*m1*(d_m_sq_sol+m1**2.)**0.5\n",
    "        m1,opt_output,success,output_message = fsolve(m1_func,sum_masses/3.,(sum_masses,delta_m_squared_atm,delta_m_squared_sol),full_output=True)\n",
    "        m1 = m1[0]\n",
    "        m2 = (delta_m_squared_sol + m1**2.)**0.5\n",
    "        m3 = (delta_m_squared_atm + 0.5*(m2**2. + m1**2.))**0.5\n",
    "        return m1,m2,m3\n",
    "    else:\n",
    "        # Inverted hierarchy massive neutrinos. Calculates the individual\n",
    "        # neutrino masses from M_tot_IH and deletes M_tot_IH\n",
    "        #delta_m_squared_atm=-2.45e-3\n",
    "        #delta_m_squared_sol=7.50e-5\n",
    "        delta_m_squared_atm = -delta_m_squared_atm\n",
    "        m1_func = lambda m1, M_tot, d_m_sq_atm, d_m_sq_sol: M_tot**2. + 0.5*d_m_sq_sol - d_m_sq_atm + m1**2. - 2.*M_tot*m1 - 2.*M_tot*(d_m_sq_sol+m1**2.)**0.5 + 2.*m1*(d_m_sq_sol+m1**2.)**0.5\n",
    "        m1,opt_output,success,output_message = fsolve(m1_func,sum_masses/3.,(sum_masses,delta_m_squared_atm,delta_m_squared_sol),full_output=True)\n",
    "        m1 = m1[0]\n",
    "        m2 = (delta_m_squared_sol + m1**2.)**0.5\n",
    "        m3 = (delta_m_squared_atm + 0.5*(m2**2. + m1**2.))**0.5\n",
    "        return m1,m2,m3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test of this function, returning the 3 masses for total mass of 0.1eV\n",
    "m1,m2,m3 = get_masses(2.45e-3,7.50e-5,0.1,'NH')\n",
    "print ('NH:',m1,m2,m3,m1+m2+m3)\n",
    "m1,m2,m3 = get_masses(2.45e-3,7.50e-5,0.1,'IH')\n",
    "print ('IH:',m1,m2,m3,m1+m2+m3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The goal of this cell is to compute the ratio of P(k) for NH and IH with the same total mass\n",
    "commonsettings = {'N_ur':0,\n",
    "                  'N_ncdm':3,\n",
    "                  'output':'mPk',\n",
    "                  'P_k_max_1/Mpc':3.0,\n",
    "                  # The next line should be uncommented for higher precision (but significantly slower running)\n",
    "                  'ncdm_fluid_approximation':3,\n",
    "                  # You may uncomment this line to get more info on the ncdm sector from Class:\n",
    "                  'background_verbose':1\n",
    "                 }\n",
    "\n",
    "# array of k values in 1/Mpc\n",
    "kvec = np.logspace(-4,np.log10(3),100)\n",
    "# array for storing legend\n",
    "legarray = []\n",
    "\n",
    "# loop over total mass values\n",
    "for sum_masses in [0.1, 0.115, 0.13]:\n",
    "    # normal hierarchy\n",
    "    [m1, m2, m3] = get_masses(2.45e-3,7.50e-5, sum_masses, 'NH')\n",
    "    NH = Class()\n",
    "    NH.set(commonsettings)\n",
    "    NH.set({'m_ncdm':str(m1)+','+str(m2)+','+str(m3)})\n",
    "    NH.compute()\n",
    "    # inverted hierarchy\n",
    "    [m1, m2, m3] = get_masses(2.45e-3,7.50e-5, sum_masses, 'IH')\n",
    "    IH = Class()\n",
    "    IH.set(commonsettings)\n",
    "    IH.set({'m_ncdm':str(m1)+','+str(m2)+','+str(m3)})\n",
    "    IH.compute()\n",
    "    pkNH = []\n",
    "    pkIH = []\n",
    "    for k in kvec:\n",
    "        pkNH.append(NH.pk(k,0.))\n",
    "        pkIH.append(IH.pk(k,0.))\n",
    "    NH.struct_cleanup()\n",
    "    IH.struct_cleanup()\n",
    "    # extract h value to convert k from 1/Mpc to h/Mpc\n",
    "    h = NH.h()\n",
    "    plt.semilogx(kvec/h,1-np.array(pkNH)/np.array(pkIH))\n",
    "    legarray.append(r'$\\Sigma m_i = '+str(sum_masses)+'$eV')\n",
    "plt.axhline(0,color='k')\n",
    "plt.xlim(kvec[0]/h,kvec[-1]/h)\n",
    "plt.xlabel(r'$k [h \\mathrm{Mpc}^{-1}]$')\n",
    "plt.ylabel(r'$1-P(k)^\\mathrm{NH}/P(k)^\\mathrm{IH}$')\n",
    "plt.legend(legarray)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.savefig('neutrinohierarchy.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
